package com.hk.core.web.reactive;

import io.netty.buffer.Unpooled;
import lombok.AccessLevel;
import lombok.NoArgsConstructor;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.core.io.buffer.DefaultDataBuffer;
import org.springframework.core.io.buffer.NettyDataBuffer;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpRequestDecorator;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.util.Objects;
import java.util.function.Function;

/**
 * 参考 spring cloud gateway 中的此类
 *
 * @author Kevin
 * @date 2022/7/5 23:30
 */
@NoArgsConstructor(access = AccessLevel.PRIVATE)
public abstract class ServerWebExchangeUtils {

    public static final String CACHED_REQUEST_BODY_ATTR = "cachedRequestBody";

    private static final byte[] EMPTY_BYTES = {};

    /**
     * 缓存request Body,可多次读取
     */
    public static <T> Mono<T> cacheRequestBody(ServerWebExchange exchange, Function<ServerHttpRequest, Mono<T>> function) {
        var response = exchange.getResponse();
        var factory = response.bufferFactory();
        return DataBufferUtils.join(exchange.getRequest().getBody())
                .defaultIfEmpty(factory.wrap(EMPTY_BYTES))
                .map(dataBuffer -> decorate(exchange, dataBuffer))
                .switchIfEmpty(Mono.just(exchange.getRequest()))
                .flatMap(function);
    }

    private static ServerHttpRequest decorate(ServerWebExchange exchange, DataBuffer dataBuffer) {
        if (dataBuffer.readableByteCount() > 0) {
            exchange.getAttributes().put(CACHED_REQUEST_BODY_ATTR, dataBuffer);
        }
        return new ServerHttpRequestDecorator(exchange.getRequest()) {
            @Override
            public Flux<DataBuffer> getBody() {
                return Mono.fromSupplier(() -> {
                    if (Objects.isNull(exchange.getAttribute(CACHED_REQUEST_BODY_ATTR))) {
                        return null;
                    }
                    if (dataBuffer instanceof NettyDataBuffer pdb) {
                        return pdb.factory().wrap(pdb.getNativeBuffer().retainedSlice());
                    } else if (dataBuffer instanceof DefaultDataBuffer ddf) {
                        return ddf.factory().wrap(Unpooled.wrappedBuffer(ddf.getNativeBuffer()).nioBuffer());
                    } else {
                        throw new IllegalArgumentException(
                                STR."Unable to handle DataBuffer of type \{dataBuffer.getClass()}");
                    }
                }).flux();
            }
        };
    }
}
